'''
Version: 1.0
Author: hpf
Date: 2021-05-28 01:44:14
LastEditTime: 2021-07-30 04:23:00
LastEditors: hpf
Description: 检测算法运行进程
FilePath: /trackrunner/detectrunner.py
'''
import importlib
import json
import os
import sys

import cv2
import yaml
'''
@description: 视频处理接口
@param {*} library_path 运行的算法路径
@param {*} model_path 运行算法的模型路径
@param {*} config_path 运行算法的配置文件路径
@param {*} input_path 输入算法处理的视频路径
@param {*} output_path 输出结果的文件夹
@return {*}
'''


def run_video(library_path, model_path, config_path, input_path, output_path):
    print('Starting run video...lib:{},model:{},config:{},input:{},output:{}'.
          format(library_path, model_path, config_path, input_path,
                 output_path))

    process_path = sys.argv[0]
    sys.argv.clear()
    sys.argv.append(process_path)

    # 加载配置文件
    names_list = []
    with open('./coco.names') as names_file:
        for line in names_file.readlines():
            names_list.append(line.strip('\n'))

    # 导入算法库
    os.chdir(library_path)
    sys.path.append(library_path)
    object = importlib.import_module("tools.detector")

    if config_path == "":
        config_path = library_path + "/detect_cfg.yaml"

    with open(config_path) as f:
        cfg_yaml = yaml.load(f, Loader=yaml.FullLoader)  # model dict

    if model_path != "":
        cfg_yaml['TEST']['weight'] = model_path
    cfg_yaml['GeneralOptions']['type'] = "TEST"

    # 初始化变量
    vc = cv2.VideoCapture(input_path)
    frame_id = 0
    video_count = vc.get(cv2.CAP_PROP_FRAME_COUNT)
    process_num = -1

    # 初始化检测算法
    detector = object.ZJLAB_DETECT(cfg_yaml)

    while (vc.isOpened()):
        ret, frame = vc.read()
        if int(100 * frame_id / video_count) != process_num:
            process_num = int(100 * frame_id / video_count)
            print("PlatformProcessNum:%d" % (process_num))
        if ret:
            # 运行run处理图片
            results = detector.run(frame)

            # 保存结果
            for item in results:
                try:
                    cv2.rectangle(frame, (int(
                    item[0] * frame.shape[1]), int(item[1] * frame.shape[0])),
                              (int(item[2] * frame.shape[1]),
                               int(item[3] * frame.shape[0])), (255, 0, 0), 2)
                    cv2.putText(frame, names_list[item[5]] + "|" + str(item[4]),
                            (int(item[0] * frame.shape[1]),
                             int(item[1] * frame.shape[0])),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
                except:
                    print("error")
                
            cv2.imwrite(os.path.join(output_path, "%010d.jpg" % (frame_id)),
                        frame)

            frame_id += 1
        else:
            break
    print("PlatformProcessNum:100")


'''
@description: 数据集处理接口
@param {*} library_path 运行的算法路径
@param {*} model_path 运行算法的模型路径
@param {*} config_path 运行算法的配置文件路径
@param {*} input_path 输入算法处理的数据集路径
@param {*} output_path 输出结果的文件夹
@return {*}
'''


def run_eval(library_path, model_path, config_path, input_path, output_path):
    print('Starting run eval...lib:{},model:{},config:{},input:{},output:{}'.
          format(library_path, model_path, config_path, input_path,
                 output_path))

    process_path = sys.argv[0]
    sys.argv.clear()
    sys.argv.append(process_path)

    # 加载配置文件
    names_list = []
    with open('./coco.names') as names_file:
        for line in names_file.readlines():
            names_list.append(line.strip('\n'))

    # 导入算法库
    os.chdir(library_path)
    sys.path.append(library_path)
    object = importlib.import_module("tools.detector")

    if config_path == "":
        config_path = library_path + "/detect_cfg.yaml"

    with open(config_path) as f:
        cfg_yaml = yaml.load(f, Loader=yaml.FullLoader)
    file_path = os.path.join(input_path, "coco/images/val2017/")

    if model_path != "":
        cfg_yaml['TEST']['weight'] = model_path
    cfg_yaml['GeneralOptions']['type'] = "TEST"

    # 初始化变量
    file_lists = os.listdir(file_path)
    img_count = len(file_lists)
    img_id = 0
    process_num = -1

    # 初始化检测算法
    detector = object.ZJLAB_DETECT(cfg_yaml)
    for i in range(0, img_count):
        img_path = os.path.join(file_path, file_lists[i])
        if int(100 * img_id / img_count / 2) != process_num:
            process_num = int(100 * img_id / img_count / 2)
            print("PlatformProcessNum:%d" % (process_num))
        frame = cv2.imread(img_path)
        if frame is None:
            continue

        # 运行run处理图片
        results = detector.run(frame)

        # 保存结果
        for item in results:
            cv2.rectangle(
                frame,
                (int(item[0] * frame.shape[1]), int(item[1] * frame.shape[0])),
                (int(item[2] * frame.shape[1]), int(item[3] * frame.shape[0])),
                (255, 0, 0), 2)
            if item[5] < len(names_list):
                try:
                    cv2.putText(frame,
                            names_list[int(item[5])] + "|" + str(item[4]),
                            (int(item[0] * frame.shape[1]),
                             int(item[1] * frame.shape[0])),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
                except:
                    print("error")
                
        cv2.imwrite(os.path.join(output_path, file_lists[i]), frame)
        img_id += 1

    # 修改配置文件
    if model_path != "":
        cfg_yaml['EVAL']['weight'] = model_path
    cfg_yaml['EVAL']['coco_path'] = input_path + "/coco"
    cfg_yaml['GeneralOptions']['type'] = "EVAL"

    # 重新初始化检测算法
    detector_eval = object.ZJLAB_DETECT(cfg_yaml)

    # 计算AP值
    stats = detector_eval.evaluate()
    dict = {
        "Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ]":
        stats[0],
        "Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ]":
        stats[1]
    }
    print("PlatformSummary:%s" % (json.dumps(dict)))
    print("PlatformProcessNum:100")


if __name__ == '__main__':
    if int(sys.argv[6]) == 1:
        run_eval(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4],
                 sys.argv[5])
    else:
        run_video(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4],
                  sys.argv[5])
